import os
import matplotlib
matplotlib.use('Qt4Agg')
import pylab as plt
import processGPR as gpr

class Viewer:
    
    def __init__(self):
        ''' Set parameters '''
        self.cmin = 122
        self.cmax = 142
        
        self.le = None
        self.ri = None
        self.procdir = 'F:\IPYdata\CBAND\\processed\\'
        self.traceNo = 0
        
        ''' Set up the figure and axes '''
        self.fig = plt.figure(None, figsize = (14, 8))
        
        self.axprofile = plt.axes([0.075, 0.075, 0.65, 0.85])
        self.axtrace = plt.axes([0.75, 0.075, 0.2, 0.85])        
        
        self.tracexlabel = 'Power [dB]'
        
        ''' Connect events '''
        self.connect()
    
    
    def connect(self):
        self.cidgetleft = self.fig.canvas.mpl_connect('key_press_event', self.getleft)
        self.cidgetright = self.fig.canvas.mpl_connect('key_press_event', self.getright)
        self.cidremove = self.fig.canvas.mpl_connect('key_press_event', self.remove)
        self.cidsave = self.fig.canvas.mpl_connect('key_press_event', self.save)
        self.cidnext = self.fig.canvas.mpl_connect('key_press_event', self.next)
        self.cidwigglewin = self.fig.canvas.mpl_connect('motion_notify_event', self.wiggle)
            
    
    def getList(self, filelist):
        self.file = iter(filelist)
        
    
    def getFile(self, filename):
        self.file = filename
        
    
    def getData(self, TD):
        self.profile = TD
        
    
    def loadData(self):
        if isinstance(self.file, list):
            file = self.file.next()
        elif isinstance(self.file, str):
            file = self.file
        else:
            file = None
            
        if file[-4:-2] == '.g':
            FD = gpr.CFD()
            FD.readRaw(file[:-4])
            self.profile = gpr.CTD()
            self.profile.timeDomain(FD.data)
            self.profile.staticCorrection()
            self.profile.wscs()
#            if len(file[7:-4]) == 1:
#                self.profile.filename = self.folder+'_0'+file[7:-4]+'.b'+file[-2:]
#            else:
#                self.profile.filename = self.folder+'_'+file[7:-4]+'.b'+file[-2:]
                
        elif file[-4:] == '.rd3':
            self.profile = gpr.IMP()
            self.profile.loadrd3(file[:-4])
            self.profile.subtractDC()
            self.cmin = -500 #self.profile.data.min(axis=None)
            self.cmax = 500 #self.profile.data.max(axis=None)
            self.tracexlabel = 'Amplitude'
            
        elif file == None:
            pass
        else:
            self.profile = p.CTD()
            self.profile.readBin(file, 805)
            self.profile.filename = file
        
        
    def draw(self):
        ''' Profile axes '''
        self.axprofile.imshow(self.profile.data, aspect='auto', cmap=plt.cm.binary,  vmin=self.cmin, vmax=self.cmax,
                              extent=[0, self.profile.traces, self.profile.t[-1], 0])
        
        self.axprofile.set_xlabel('Trace #')
        self.axprofile.set_ylabel('Time [ns]')
        self.axprofile.set_title(self.profile.filename)
        self.axprofile.set_xlabel('Trace #')
        self.axprofile.set_ylabel('Time [ns]')
        
        ''' Trace axes '''
        self.axtrace.plot(self.profile.data[:, self.traceNo], self.profile.t, color='k') 
        self.axtrace.axis([self.cmin, self.cmax, self.axprofile.axis()[2],  self.axprofile.axis()[3]])
        self.axtrace.set_title('Trace # %i' % self.traceNo)  
        self.axtrace.set_xlabel(self.tracexlabel)
        self.axtrace.set_ylabel('Time [ns]') 
        self.axtrace.yaxis.set_label_position('right')
        self.axtrace.yaxis.tick_right()
        
        self.fig.canvas.draw()     
        
        
    def show(self):
        self.loadData()
        self.draw()
        plt.show()   
        
    
    ''' KEY EVENTS '''
        
    def getleft(self, event):
        if event.inaxes:
            if event.key == 'q':
                self.le = int(event.xdata)
                print self.le
            
            
    def getright(self, event):
        if event.inaxes:
            if event.key == 'w':
                self.ri = int(event.xdata)
                print self.ri
           

    def remove(self, event):
        if event.key == 'r':
            self.profile.removeTraces([self.le, self.ri])
            self.le = None
            self.ri = None
            self.draw()

    
    def save(self, event):
        if event.key == 's':
            self.profile.writeBin(self.procdir+self.profile.filename)
    
    
    def next(self, event):
        if event.key == 'n':
            self.ax1.clear()
            self.loadData()
            self.draw()
    
    
    ''' MOTION EVENTS '''
    
    def wiggle(self, event):
        if event.inaxes == self.axprofile.axes:
            try:
                self.traceNo = int(event.xdata)
                self.axtrace.hold(False)
                self.axtrace.plot(self.profile.data[:,self.traceNo], self.profile.t, color='k') 
                self.axtrace.hold(True)  
                self.axtrace.axhline(event.ydata, color='r') 
                self.axtrace.axis([self.cmin, self.cmax, self.axprofile.axis()[2],  self.axprofile.axis()[3]])
                self.axtrace.set_title('Trace # %i' % self.traceNo)  
                self.axtrace.set_xlabel(self.tracexlabel)
                self.axtrace.set_ylabel('Time [ns]')   
                self.axtrace.yaxis.set_label_position('right')
                self.axtrace.yaxis.tick_right()     
                self.fig.canvas.draw()
            except OverflowError:
                print 'No valid data!'


if __name__ == '__main__':
    VU = Viewer()
    VU.getFile('F:\\IPYdata\\RAMAC\\071125\\Profile23.rd3')
    VU.show()
    
